﻿// TCP-Forward.cpp : Defines the entry point for the application.
//

#include "TCP-Forward.h"



#include <iostream>
#include <boost/asio.hpp>
#include <memory>
#include <iomanip>
#include <functional>

// 使用boost::asio库中的ip和系统模块
using boost::asio::ip::tcp;
using boost::system::error_code;
// 使用标准库中的占位符
using std::placeholders::_1;
using std::placeholders::_2;

/**
 * Session类负责处理客户端和服务器之间的数据转发
 */
class Session : public std::enable_shared_from_this<Session> {
public:
    // 构造函数：初始化客户端套接字、目标主机和端口
    Session(tcp::socket client_socket, const std::string& forward_host, unsigned short forward_port)
        : client_(std::move(client_socket)),
        forward_host_(forward_host),
        forward_port_(forward_port),
        server_(client_.get_executor()) {
    }

    // 启动会话，开始连接到服务器
    void start() {
        connect_to_server();
    }

private:
    // 连接到服务器的逻辑
    void connect_to_server() {
        auto self(shared_from_this());
        tcp::resolver resolver(client_.get_executor());
        tcp::resolver::results_type endpoints = resolver.resolve(forward_host_, std::to_string(forward_port_));

        async_connect(server_, endpoints,
            [this, self](const error_code& ec, const tcp::endpoint&) {
                if (!ec) {
                    print_connection_info();
                    start_relay();
                }
                else {
                    std::cerr << "Server connection error: " << ec.message() << std::endl;
                }
            });
    }

    // 开始双向转发数据
    void start_relay() {
        start_client_to_server();
        start_server_to_client();
    }

    // 从客户端读取数据并转发到服务器
    void start_client_to_server() {
        client_.async_read_some(boost::asio::buffer(client_buffer_),
            std::bind(&Session::handle_client_read, shared_from_this(), _1, _2));
    }

    // 从服务器读取数据并转发到客户端
    void start_server_to_client() {
        server_.async_read_some(boost::asio::buffer(server_buffer_),
            std::bind(&Session::handle_server_read, shared_from_this(), _1, _2));
    }

    // 处理从客户端读取数据的完成
    void handle_client_read(const error_code& ec, size_t length) {
        if (!ec) {
            print_transfer_info("Client -> Server", length);
            async_write(server_, boost::asio::buffer(client_buffer_, length),
                std::bind(&Session::handle_server_write, shared_from_this(), _1, _2));
        }
        else {
            shutdown();
        }
    }


    void handle_client_read22(const error_code& ec, size_t length) {
        if (!ec) {
            print_transfer_info("Client -> Server", length);
            auto self(shared_from_this());
            async_write(server_, boost::asio::buffer(client_buffer_, length),
                [this,self](const error_code& ec, size_t) {
                    if (!ec) {
                        start_client_to_server();
                    }
                    else {
                        shutdown();
                    }
                });
               //std::bind(&Session::handle_server_write, shared_from_this(), _1, _2));
        }
        else {
            shutdown();
        }
    }


    // 处理写入服务器数据的完成
    void handle_server_write(const error_code& ec, size_t) {
        if (!ec) {
            start_client_to_server();
        }
        else {
            shutdown();
        }
    }

    /**
     * @brief 处理从服务器端读取数据的回调函数
     *
     * 当服务器端数据到达时，此函数将被触发。根据读取结果决定数据转发或关闭连接。
     *
     * @param ec 异步操作结果错误码，用于判断读取是否成功
     * @param length 实际接收到的数据字节长度，用于数据转发操作
     * @note 无返回值
     */
    void handle_server_read(const error_code& ec, size_t length) {
        // 成功读取服务器数据时的处理流程
        if (!ec) {
            // 打印接收信息并转发数据到客户端
            print_receive_info("Server -> Client", length);

            // 异步将服务器数据写入客户端，并绑定客户端写完成回调
            async_write(client_, boost::asio::buffer(server_buffer_, length),
                std::bind(&Session::handle_client_write, shared_from_this(), _1, _2));
        }
        // 读取异常时的处理
        else {
            // 关闭客户端和服务器的连接
            shutdown();
        }
    }

    // 处理写入客户端数据的完成
    void handle_client_write(const error_code& ec, size_t) {
        if (!ec) {
            start_server_to_client();
        }
        else {
            shutdown();
        }
    }

    // 关闭套接字
    void shutdown() {
        error_code ec;
        client_.shutdown(tcp::socket::shutdown_both, ec);
        server_.shutdown(tcp::socket::shutdown_both, ec);
    }

    // 打印连接信息
    void print_connection_info() {
        std::cout << "New connection: "
            << client_.remote_endpoint() << " -> "
            << server_.remote_endpoint() << std::endl;
    }

    // 打印数据传输信息
    void print_transfer_info(const std::string& direction, size_t length) {
        std::cout << "[" << direction << "] "
            << client_.remote_endpoint() << " --> "
            << server_.remote_endpoint() << " | "
            << length << " bytes" << std::endl;
    }

    // 打印数据接收信息
    void print_receive_info(const std::string& direction, size_t length) {
        std::cout << "[" << direction << "] "
            << client_.remote_endpoint() << " <-- "
            << server_.remote_endpoint() << " | "
            << length << " bytes" << std::endl;
    }

    // 客户端的套接字
    tcp::socket client_;
    // 服务器的套接字
    tcp::socket server_;
    // 数据缓冲区
    std::array<char, 4096> client_buffer_;
    std::array<char, 4096> server_buffer_;
    // 目标主机和端口
    std::string forward_host_;
    unsigned short forward_port_;
};

/**
 * PortForwarder类负责监听指定端口，并为每个连接创建一个新的Session实例
 */
class PortForwarder {
public:
    // 构造函数：初始化监听端口、目标主机和端口
    PortForwarder(boost::asio::io_context& io, unsigned short listen_port,
        const std::string& forward_host, unsigned short forward_port)
        : acceptor_(io, tcp::endpoint(tcp::v4(), listen_port)),
        forward_host_(forward_host),
        forward_port_(forward_port)
    {
        // 监听端口
    }

    // 开始接受连接
    void start_accept() {
        acceptor_.async_accept(
            std::bind(&PortForwarder::handle_accept, this, _1, _2));
    }

private:
    // 处理接受连接的完成
    void handle_accept(const error_code& ec, tcp::socket socket) {
        if (!ec) {
            std::make_shared<Session>(std::move(socket), forward_host_, forward_port_)->start();
        }
        start_accept();
    }

    // 监听器
    tcp::acceptor acceptor_;
    // 监听器和目标主机、端口
    std::string forward_host_;
    unsigned short forward_port_;
};

// 启动端口转发功能
int startForward() {
    // 读取用户输入
    std::string input;
    std::getline(std::cin, input); // 读取整行输入

    std::istringstream iss(input);
    std::string word;
    std::vector<std::string> words;

    // 解析输入
    while (iss >> word) {
        words.push_back(word);
    }

    // 创建并启动端口转发器
    try {
        boost::asio::io_context io;
        PortForwarder forwarder(
            io,
            static_cast<unsigned short>(std::stoi(words.at(0))),
            words.at(1),
            static_cast<unsigned short>(std::stoi(words.at(2))));

        forwarder.start_accept();
        std::cout << "正在监听端口:" << words.at(0) << " 转发到" << words.at(1) << ":" << words.at(2) << std::endl;

        io.run();
        std::cout << "监听结束" << std::endl;
    }
    catch (std::exception& e) {
        std::cerr << "Exception: " << e.what() << "\n";
    }

    return 0;
}


int main()
{
	std::cout << " ***** TCP端口转发程序V1.0(2025-6-5 bst2000) ***** " << std::endl;
    std::cout << "请输入: 监听端口  转发目的IP  转发目的端口，用空格分隔" << std::endl;
    startForward();
	return 0;
}

